
import os
from svmutil import *

negative_corpus_path="C:/Users/Jorge/Documents/sentiment/data/imdb1/neg"  # insert the path to the directory of interest
positive_corpus_path="C:/Users/Jorge/Documents/sentiment/data/imdb1/pos"
negative_test_corpus_path="C:/Users/Jorge/Documents/sentiment/data/imdb1/train/neg"
positive_test_corpus_path= "C:/Users/Jorge/Documents/sentiment/data/imdb1/train/pos"

def get_files_of_directory(directory):
	files = []
	dirList=os.listdir(directory)
	for fname in dirList:
		files.append(fname)
	return files

def read_document(path):

	file = open(path)
	text = file.read()
	tokens = text.split()
	return tokens


def create_vocabulary(list_directory):

	vocabulary = {}

	for directory in list_directory:
		files = get_files_of_directory(directory)
		for file in files:
			tokens = read_document(directory+'/'+file)
			for token in tokens:
				try:
					vocabulary[token] +=1
				except KeyError:
					vocabulary[token] = 1

	print len(vocabulary)
	vocabulary = reduce_vocabulary(vocabulary, len(vocabulary)*0.6)
	return vocabulary


def reduce_vocabulary(vocabulary, new_size):

	values = vocabulary.values()
	values.sort()
	values.reverse()
	if (len(vocabulary)<new_size):
		return None
	values =  values[0: int(new_size)]
	new_vocabulary = {}
	for word in vocabulary.keys():
		if vocabulary[word] in values:
			new_vocabulary[word] = vocabulary[word]

	return new_vocabulary



def vectorizer_document(vocabulary , tokens):

	vector = range(len(vocabulary))
	for i in range(len(vocabulary)):
		if vocabulary[i] in tokens:
			vector[i] = 1
		else:
			vector[i] = 0

	return vector

def get_documents_from_path(path):

	documents = []
	for file in get_files_of_directory(path):
		tokens = read_document(path+'/'+file)
		documents.append(tokens)
	return documents

#tokens = read_file('C:/Users/Jorge/Documents/sentiment/data/imdb1/neg/cv000_29416.txt')
#print len(tokens)
print 'creando vocabulario...'
vocabulary = create_vocabulary([negative_corpus_path, positive_corpus_path])
print len(vocabulary)
#for key in vocabulary.keys():
#	print key +' -> '+ str(vocabulary[key])
print 'creando datos de entrenamiento'
X = []
Y = []

print 'vectorizando documentos negativos'
for doc in get_documents_from_path(negative_corpus_path):
	v = vectorizer_document(vocabulary.keys(), doc)
	X.append(v)

Y = [0]*len(X)

print 'vectorizando documentos positivos' 
for doc in get_documents_from_path(positive_corpus_path):
	v = vectorizer_document(vocabulary.keys(), doc)
	X.append(v) 

Y.extend( [1]*(len(X)-len(Y)) )

print 'entrenando svm...'

model = svm_train(Y, X)

print 'probando..'

test_X = []
test_Y = []
for doc in get_documents_from_path(negative_test_corpus_path):
	v = vectorizer_document(vocabulary.keys(), doc)
	test_X.append(v) 

test_Y = [0]*len(test_X)

for doc in get_documents_from_path(positive_test_corpus_path):
	v = vectorizer_document(vocabulary.keys(), doc)
	test_X.append(v)

test_Y.extend( [1]*(len(test_X)-len(test_Y)) )


p_labels, p_acc, p_vals =  svm_predict(test_Y, test_X, model)
print p_labels
print p_acc
print p_vals
